package com.sxu.search.bloom;

import java.util.BitSet;
import java.util.Random;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;

/**
 * 布隆过滤器
 */
public class BloomFilter {
    private BitSet bitSet; // 位数组
    private int bitSize; // 位数组长度
    private int hashFuncCount; // hash函数个数
    private int[] seeds; // 固定的种子池 确保添加和查询时种子相同

    // 传入预期存储个数 and 允许的误判率
    public BloomFilter(int expectSize, double falsePositiveRate) {
        this.bitSize = calculateBitSize(expectSize, falsePositiveRate);
        this.hashFuncCount = calculateHashFuncCount(bitSize,expectSize);
        this.bitSet = new BitSet(bitSize);
        this.seeds = getSeeds(hashFuncCount);
    }

    // 根据误判率公式计算m和k
    // 根据期望的容量和错误率计算位数组大小
    private int calculateBitSize(int n, double p) {
        return (int) (-n * Math.log(p) / (Math.log(2) * Math.log(2)));
    }

    // 计算hash函数个数
    private int calculateHashFuncCount(int m, int n) {
        return Math.max(1, (int) (m / n * Math.log(2)));
    }

    // 生成种子池
    private int[] getSeeds(int k){
        int[] seeds = new int[k];
        Random random = new Random();
        for(int i=0;i<k;i++){
            seeds[i] = random.nextInt();
        }
        return seeds;
    }

    // 计算hash值
    private int hash(byte[] data, int seed) {
        try {
            // 使用MD5生成哈希值
            MessageDigest md = MessageDigest.getInstance("MD5");
            md.update((byte) seed);      // 添加随机种子
            byte[] digest = md.digest(data);
            return Math.abs(new String(digest).hashCode());
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException("MD5 not available");
        }
    }

    // 添加元素
    public void add(String element) {
        byte[] bytes = element.getBytes();
        for(int i =0;i<this.hashFuncCount;i++) {
            // 每次计算使用不同的随机值 模拟多个hash函数
            int hash = hash(bytes,seeds[i]);
            bitSet.set(hash % this.bitSize,true); // 取模确保在合法范围内
        }
    }

    // 查询元素
    public boolean contains(String element) {
        byte[] bytes = element.getBytes();
        for(int i =0;i<this.hashFuncCount;i++) {
            int hash = hash(bytes,seeds[i]);
            if(!bitSet.get(hash % this.bitSize)) {
                return false; // 一定不存在
            }
        }
        return true; // 可能存在
    }
}
